from src.module.TongyiModel import TongyiModel
from src.constance import graph_image

from langchain.tools import tool
from langgraph.prebuilt import create_react_agent, ToolNode
from langgraph.graph import StateGraph, START, END, MessagesState
from langgraph.checkpoint.memory import MemorySaver
# from langgraph.checkpoint.sqlite import SqliteSaver ## 生产环境使用 sqllit 进行持久化存储

"""
    langgraph 本质上是一个流程，流程中存在很多节点，这些节点都存在先后顺序。
    想要完成一个流程就需要让这些节点按顺序工作。每个节点可以返回一些数据，当数据返回后，可以给到后续的节点处理，因而最终得到一个结果，这就是 langgraph的本质。
    因此，使用langgraph时，不一定要让ai嵌入，ai的嵌入是在：一个节点中存在很多工具时，让ai去推理选择一个合适的工具来处理回复。
"""

##  工具函数
class Tools:
    @staticmethod
    @tool
    def search_weather(address: str):
        """ 查询城市天气

        Args:
            address: 地址
        """
        if address == '深圳':
            return '晴朗'
        else:
            return '多云'

    @staticmethod
    @tool
    def search_road(address: str):
        """ 查询城市交通情况

        Args:
            address : 地址
        """
        if address == '深圳':
            return '道路通畅'
        else:
            return  '拥堵'

## 基础简单示例
class BaseSimple(TongyiModel):
    def start(self):
        agent = create_react_agent(model=self.model, tools=[Tools().search_weather])
        response = agent.invoke({ "messages": "我现在即将抵达省份: 湖北省，城市：襄阳市枣阳市，城镇：鹿头镇，天气怎么样?" })
        print(response)

## 进阶版1
class Advanced1:
    def __init__(self):
        self.model = TongyiModel().model

        # 将工具作为一个节点，并绑定在大模型上面， 因为工具本身不能作为节点，因此需要 ToolNode 辅助函数来实现
        tools = [Tools().search_road, Tools().search_weather]
        self.node_tools = ToolNode(tools)
        self.model = self.model.bind_tools(tools)

    # 将节点流向转换为图片进行保存
    @staticmethod
    def draw_flow(graph, filename: str):
        flow_image = graph.get_graph().draw_mermaid_png()
        with open(str(graph_image(filename)), 'wb') as f:
            f.write(flow_image)

    # 聊天节点
    def chatbot(self, state: MessagesState):
        return { "messages": [self.model.invoke(state['messages'])] }

    # 条件边，判断下一步的走向
    @staticmethod
    def conditional(state: MessagesState):
        messages = state['messages'][-1]
        if messages.tool_calls:
            return 'tool'
        return END

    @staticmethod
    def get_message_content(response):
        return response['messages'][-1].content

    def start(self):
        graph_builder = StateGraph(MessagesState)
        graph_builder.add_node('chatbot', self.chatbot) # 设置 chatbot 节点
        graph_builder.add_node('tool', self.node_tools) # 将工具作为节点绑定在 graph 上面。


        graph_builder.add_edge(START, 'chatbot') # 普通边， start 到 chatbot
        graph_builder.add_conditional_edges('chatbot', self.conditional) # 条件边，判断是否需要走到下一步 chatbot 到 tool
        graph_builder.add_edge('tool', 'chatbot') # 普通边，tool 结果返回 到 chatbot
        graph_builder.add_edge('chatbot', END) # 普通边，chatbot 到 end 结束

        # 添加消息记录
        pointer_memo_saver = MemorySaver()
        # pointer_memo_saver = SqliteSaver()
        graph = graph_builder.compile(checkpointer=pointer_memo_saver)

        # 绘制流程图
        self.draw_flow(graph, 'advanced1.png')

        response = graph.invoke({"messages": "深圳的天气怎么样? "}, config={'configurable': { 'thread_id': 123 }})
        print(self.get_message_content(response))

        response_city = graph.invoke({"messages": "我问的是哪个城市?"}, config={'configurable': {'thread_id': 123}})
        print(self.get_message_content(response_city))

class LangGraphExample:
    def __init__(self):
        # self.base_simple = BaseSimple()  # 基础简单示例
        self.advanced1 = Advanced1() # 进阶示例1


    def start(self):
        # self.base_simple.start()
        self.advanced1.start()
